# -*- coding: utf-8 -*-

"""
@Author      : Sky Bully
@Time        : 2024/6/6 10:48
"""
import sys
import os
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import matplotlib.pyplot as plt

# 字体加载
from matplotlib import font_manager, rcParams

font_path = r"D:\Applications\字体合并\字体合并补全工具-简体中文-1.1.0-windows-x64\out (2).ttf"
font_manager.fontManager.addfont(font_path)
prop = font_manager.FontProperties(fname=font_path)

# 字体设置
rcParams['font.family'] = 'sans-serif'  # 使用字体中的无衬线体
rcParams['font.sans-serif'] = prop.get_name()  # 根据名称设置字体
rcParams['font.size'] = 10  # 设置字体大小
rcParams['axes.unicode_minus'] = False  # 使坐标轴刻度标签正常显示正负号


labels = {'CH3O2': 'CH$_{3}$O$_{2}$(ppb)', 'NO3': 'NO$_{3}$(ppb)', 'OH': 'OH(ppb)', 'HO2': 'HO$_{2}$(ppb)',
          'O3': 'O$_{3}$(ppb)'}


def main_analysis(species_file, start_date, factors, out_loc):
    """

    :param factors: 绘图物种，逗号分隔，仅支持两种及两种以内数量的物种
    :param species_file: speciesConcentrations.output文件全路径
    :param start_date: 约束数据开始时间 YYYYMMDDHH
    :return:
    """
    df = pd.read_csv(species_file, sep='\\s+')
    df['t'] = df['t'].apply(lambda x: datetime.strptime(start_date, "%Y%m%d%H") + timedelta(hours=round(x / 3600 + 0)))
    df[df.columns[1:]] = df[df.columns[1:]].apply(lambda x: x / 2.46e10)

    factor_list = factors.split(',')
    for factor in factor_list:
        if factor.upper() not in df.columns:
            print(f"[ERROR]: 物种{factor}未出现在模拟结果中，请检查outputSpecies.config中是否添加改物种")
            sys.exit(-999)

    ticks = pd.date_range(start=start_date[:8], periods=len(df['t']) + int(start_date[8:]), freq='h')[int(start_date[8:]):][::4]
    tick_labels = [i.strftime("%#m月%#d日%#H时") for i in ticks]

    fig = plt.figure(figsize=(9, 2))
    ax1 = fig.add_subplot(111)
    ax1.plot(df[factor_list[0]].astype(float).values, color='#ff000d', label=labels[factor_list[0]], linewidth=2)
    plt.ylabel(f"{labels[factor_list[0]]}", fontsize=14)
    plt.legend(loc='upper left', fontsize=10, ncol=2)

    if len(factor_list) == 2:
        ax2 = ax1.twinx()
        ax2.plot(df[factor_list[1]].astype(float).values, color='#02ab2e', label=labels[factor_list[1]], linewidth=2)
        plt.ylabel(f"{labels[factor_list[1]]}", fontsize=14)
        plt.legend(loc='upper right', fontsize=10, ncol=2)

    plt.margins(x=-0.01)
    plt.xlim(-0.5, len(df['t'].index) - 0.5)
    plt.xticks(np.arange(0, len(df['t']))[::4], tick_labels, rotation=20, fontsize=14)

    if len(factor_list) >= 3:
        print(f"[Warning]: 物种{factor_list[2]}及后面的部分未参与绘图（一次性绘图仅支持两种及以内数量的物种）")

    os.makedirs(out_loc, exist_ok=True)
    plt.savefig(os.path.join(out_loc, '_'.join(factor_list) + '浓度时序图.png'), bbox_inches='tight', dpi=600)


if __name__ == '__main__':
    # if len(sys.argv) < 5:
    #     print(f"[Error]: Run with python {sys.argv[0]} './speciesConcentrations.output' '2023090108' 'CH3O2,NO3' './picture")
    # else:
    #     main_analysis(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
    main_analysis(r'E:\临时文件\OBM-MCM\speciesConcentrations (2).output', '2023090108', 'O3,NO3', './png')
